import torch
import torch.nn as nn
import matplotlib.pyplot as plt
import numpy as np

def calculate_relative_change(initial, current, eps=1e-7):
    delta = torch.abs(current - initial)
    denominator = torch.abs(initial) + eps
    return (delta / denominator).mean().item()

def getParam(model:nn.Module):
    param={name: param.data.clone() for name, param in model.named_parameters() if "fc" not in name}
    return param

def visualize_layer_changes(
    change_rates,        # 必须参数：已排序的相对变化率列表
    layer_names,         # 必须参数：对应层名称列表（与change_rates同长度）
    # title="神经网络层参数变化分析",  # 可选参数
    xlabel="layer name", 
    ylabel="relative change",
    color='#2ca02c',     # 折线主色
    marker='o',          # 数据点标记
    show_extremes=True,  # 是否标记极值点
    log_scale=True,      # 是否使用对数坐标
    save_path=None       # 图片保存路径
):
    """
    可视化神经网络各层参数变化的折线图函数
    参数说明参考：网页6[6](@ref)的封装思路 + 网页3[3](@ref)的样式配置
    """
    # 参数校验
    assert len(change_rates) == len(layer_names), "参数列表长度不一致"
    
    # 创建画布（参考网页3[3](@ref)的初始化配置）
    plt.figure(figsize=(18, 9), dpi=120)
    plt.rcParams['font.sans-serif'] = ['SimHei']
    plt.rcParams['axes.unicode_minus'] = False
    
    # 绘制折线（整合网页1[1](@ref)的plot参数配置）
    line = plt.plot(
        range(len(change_rates)),
        change_rates,
        marker=marker,
        markersize=8,
        linestyle='--',
        linewidth=2,
        color=color,
        alpha=0.7,
        label='mean relative change rate of parameter'
    )
    
    # 标记极值点（改进网页6[6](@ref)的极值标注逻辑）
    if show_extremes:
        max_idx = np.argmax(change_rates)
        min_idx = np.argmin(change_rates)
        
        plt.scatter(
            max_idx, change_rates[max_idx],
            s=150, zorder=5, edgecolors='red', facecolors='none',
            linewidths=1, label=f'changed most layer: {layer_names[max_idx]}'
        )
        plt.scatter(
            min_idx, change_rates[min_idx],
            s=150, zorder=5, edgecolors='blue', facecolors='none',
            linewidths=1, label=f'changed least layer: {layer_names[min_idx]}'
        )
    
    # 坐标轴设置（参考网页5[5](@ref)的学术图表规范）
    plt.xticks(
        range(len(layer_names)),
        layer_names,
        rotation=45,
        ha='right',
        fontsize=10
    )
    plt.yticks(fontsize=10)
    if log_scale:
        plt.yscale('log')
    plt.xlabel(xlabel, fontsize=10)
    plt.xticks(fontsize=6) 
    plt.ylabel(ylabel + ('log' if log_scale else ''), fontsize=12)
    # plt.title(title, fontsize=14, pad=20)
    
    # 辅助元素（参考网页2[2](@ref)的辅助线配置）
    plt.grid(True, linestyle=':', alpha=0.5)
    plt.legend(loc='upper left', fontsize=10)
    plt.tight_layout()
    
    # 保存或显示（整合网页7[7](@ref)的保存逻辑）
    if save_path:
        plt.savefig(save_path, bbox_inches='tight', dpi=300)
        plt.close()
        print(f"图表已保存至：{save_path}")
    else:
        plt.show()

if __name__=="__main__":
    model0=torch.load("model/trained0/best.pt",weights_only=False)
    model1=torch.load("model/trained/best.pt",weights_only=False)
    param0=getParam(model0)
    param1=getParam(model1)
    res_dict={}
    for name,param in param0.items():
        initial_param=param
        current_param=param1[name]
        rcr = calculate_relative_change(initial_param, current_param)
        res_dict[name]=rcr
    visualize_layer_changes(list(res_dict.values()),list(res_dict.keys()),save_path="compareParam.png")
    # res_list=list(res_dict.values())
    # res_list.sort(reverse=True)
    # print(res_list)